#!/usr/bin/env python
# coding: utf-8

# In[1]:


import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import os
import fenge
def predect_result(path):
    result = ''
    pipeline = transforms.Compose([
        transforms.Grayscale(num_output_channels=1),
        transforms.Resize((28, 28)),
        transforms.ToTensor(),# 将图片转换成tensor
        transforms.Normalize((0.1307,),(0.3081,)) # 正则化降低模型复杂度
    ])
    # 加载测试数据集
    folder = 'pre/class1' # 替换成你的目标文件夹路径
    for file_name in os.listdir(folder):
        file_path = os.path.join(folder, file_name)
        if os.path.isfile(file_path):
            os.remove(file_path)
    fenge.zifu_segment(path)
    test_dataset = datasets.ImageFolder('pre', transform=pipeline)
    # 创建测试数据加载器
    class Digit(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv1 = nn.Conv2d(1, 14, 5)  # 1:灰度图片的通道， 10：输出通道， 5：kernel 5x5
            self.conv2 = nn.Conv2d(14, 20, 3)  # 10:输入通道， 20：输出通道， 3：kernel 3x3
            self.fc1 = nn.Linear(2000, 500)  # 20*10*10:输入通道， 500：输出通道
            self.fc2 = nn.Linear(500, 14)  # 500：输入通道， 10：输出通道
        # 前向传播
        def forward(self, x):
            input_size = x.size(0)  # batch_size
            x = self.conv1(x)  # 输入：batch*1*28*28, 输出：batch*10*24*24 (28 - 5 + 1 = 24)
            x = F.relu(x)  # 激活函数，保持shape不变， 输出batch*10*24*24
            x = F.max_pool2d(x, 2, 2)  # 输入：batch*10*24*24 输出：batch*10*12*12
            x = self.conv2(x)  # 输入：batch*10*12*12, 输出：batch*20*10*10 (12 - 3 + 1 = 10)
            x = F.relu(x)
            x = x.view(input_size, -1)  # 拉平， -1 自动计算维度，20*10*10 = 2000
            x = self.fc1(x)  # 输入：batch*2000，输出：batch*500
            x = F.relu(x)
            x = self.fc2(x)  # 输入：batch*500，输出：batch*10
            output = F.log_softmax(x, dim=1)  # 计算分类后，每个数字的概率值
            return output
    BATCH_SIZE = 16  # 每批处理的数据
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 是否用GPU还是CPU训练
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=16, shuffle=False)
    saved_model_weights = torch.load('trained_model1.pth', map_location=torch.device('cpu'))
    #saved_model_weights = torch.load('trained_model.pth')
    # # 检查是否有CUDA支持
    # if torch.cuda.is_available():
    #     # 加载模型时将存储映射到CPU
    #     saved_model_weights = torch.load('trained_model.pth', map_location=torch.device('cpu'))
    # else:
    #     # 正常加载模型
    #     saved_model_weights = torch.load('trained_model.pth')
    # 创建模型实例并加载状态字典
    model = Digit().to(DEVICE)
    model.load_state_dict(saved_model_weights)
    # 正确率
    correct = 0.0
    # 测试损失
    test_loss = 0.0
    with torch.no_grad(): # 不会计算梯度，也不会反向传播
        for data, target in test_loader:
        # 部署到device上
            data, target = data.to(DEVICE), target.to(DEVICE)
            # 测试数据
            output = model(data)
            # 计算测试损失
            test_loss += F.cross_entropy(output, target).item()
            # 找到概率值最大的下标
            pred = output.max(1, keepdim=True)[1] # 值，索引
            # pred = torch.max(output, dim=1)
            # pred = output.argmax(dim=1)
            #累计正确的值
            correct += pred.eq(target.view_as(pred)).sum().item()
            for i in range(len(data)):
                #plt.imshow(data[i].permute(1, 2, 0).cpu())
                #plt.title(f'Predicted label: {pred[i].item()}')
                if pred[i].item()<=9:
                    result=result+str(pred[i].item())
                elif pred[i].item()==10:
                    result=result+'/'
                elif pred[i].item()==11:
                    result=result+'*'
                elif pred[i].item()==12:
                    result=result+'+'
                elif pred[i].item()==13:
                    result=result+'-'
                plt.show()
    return result


# In[ ]:




